#!/usr/bin/python2

##
## Copyright (C) 2011 Luciano Xumerle <luciano.xumerle@gmail.com>
##
## This program is free software: you can redistribute it and/or modify
## it under the terms of the GNU General Public License as published by
## the Free Software Foundation, either version 3 of the License, or
## (at your option) any later version.
##
## This program is distributed in the hope that it will be useful,
## but WITHOUT ANY WARRANTY; without even the implied warranty of
## MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
## GNU General Public License for more details.
##
## You should have received a copy of the GNU General Public License
## along with this program.  If not, see <http://www.gnu.org/licenses/>.
##

from expression import *
from knowngene import *

import time
import sys
import random
import os


def printError():
   print( "SINTAX ERROR!!!" )
   print( "Create a sequence for each row in knowngene file:" )
   print( "> " + sys.argv[0] + " <chromDir> <knownGene.txt>\n" )
   print( "Open fasta sequnce and save the .fa.bin file!" )
   print( "> " + sys.argv[0] + " -bin <chr.fa>\n" )
   print( "Get a subsequence from the .fa.bin file:" )
   print( "> " + sys.argv[0] + "sequence.fa.bin <start> <length>" )


def splitExon( ss ):
   tt = ss.split(",")
   tt.pop()
   res=[]
   for a in tt:
       res.append(int(a))
   return res


def retrieveSequence( chromosomesNames, chrom ):
   for a in chromosomesNames:
       if chrom + ".fa.bin" in a:
           return a
   return ""


def createSequence( gene, name, fragments ):
   sequence=catSequences( fragments )
   out=open( gene + '.fa.bin','w')
   if name != "":
       out.write( ">" + name + "\n")
   out.write( sequence )
   if name != "":
       out.write( "\n" )
   out.close()


def main():

   if len(sys.argv) <3:
       printError()
       return

   if sys.argv[1] == "-bin" :
       createChromosomeSequence( sys.argv[2] )
       return

   if ".fa.bin" in sys.argv[1] and len(sys.argv)==4:
      print getSequenceChromosome( sys.argv[1], int(sys.argv[2])-1, int(sys.argv[3]) )
      return

   chromFiles=dict()
   chromosomesNames=[]
   for filename in os.listdir( sys.argv[1] ):
       if ".fa.bin" in filename:
           tt = sys.argv[1] + os.sep + filename
           chromFiles[ tt ] = [ os.path.getsize( tt ) ]
           chromosomesNames.append(tt)

   if ".gtf" in sys.argv[2]:
       fragments=[]
       gene=""
       chrom=""
       sense=""
       for line in open( sys.argv[2] ):
           sequence=""
           dd = line.strip().split("\"")
           nn = dd[-2]
           if nn != gene and len(fragments)>0:
               createSequence( gene, "", fragments )
               fragments=[]
               gene=nn
           dd = line.strip().split("\t")
           if dd[2] != "exon":
               continue
           chrom = dd[0]
           chfile= retrieveSequence( chromosomesNames, chrom )
           start = int(dd[3])-1
           stop  = int(dd[4])
           sense = dd[6]
           fragments.append( [ gene, chfile, sense, start, stop ] )
       if len(fragments)>0:
           createSequence( gene, "", fragments )
   else:
       if "known" in sys.argv[2]:
           for line in open( sys.argv[2] ):
               fragments=[]
               dd = line.strip().split("\t")
               chfile=retrieveSequence( chromosomesNames, dd[1] )
               starts = splitExon(dd[8])
               stops =  splitExon(dd[9])
               for ii in range(len(starts)):
                   fragments.append( [ dd[0], chfile, dd[2], starts[ii], stops[ii] ] )
               createSequence( dd[0], "", fragments )
       else:
           ### lines in junction file are: name chr1 sense1 pos1 chr2 sense2 pos2 fragments_length
           if ".jun" in sys.argv[2]:
               for line in open( sys.argv[2] ):
                   fragments=[]
                   dd = line.strip().split("\t")
                   name=dd[0]
                   chfile1 = retrieveSequence( chromosomesNames, dd[1] )
                   chfile2 = retrieveSequence( chromosomesNames, dd[4] )
                   sense1 = dd[2]
                   sense2 = dd[5]
                   stop1 = int(dd[3])
                   start1 = stop1 - int(dd[7])
                   start2 = int(dd[6])
                   stop2 = start2 + int(dd[7])
                   fragments.append( [ name, chfile1, sense1, start1, stop1 ] )
                   fragments.append( [ name, chfile2, sense2, start2, stop2 ] )
                   createSequence( sys.argv[2], name, fragments )
           else:
               ## Each row in file is CHR SENSE START STOP (gtf coords)
               gene=sys.argv[2][ : len(sys.argv[2])-4 ]
               fragments=[]
               for line in open( sys.argv[2] ):
                   dd = line.strip().split("\t")
                   chfile=retrieveSequence( chromosomesNames, dd[0] )
                   fragments.append( [ gene, chfile, dd[1], int(dd[2])-1, int(dd[3]) ] )
               createSequence( gene, "", fragments )


# start main
if __name__=='__main__': main()
